"""
========================
气象相关的计算工具函数
从2023年中开始编写，逐渐缝缝补补
------------------------
Author: Lilidream
Contact: lilidream@outlook.com
Date: 2024-06-30
========================
"""

from xarray import DataArray, Dataset
from numpy import ndarray
import numpy as np
from . import const
from typing import Union

def butter_band_pass(data: Union[ndarray, DataArray], 
                     Tmin: float, Tmax: float , N: int = 3, fs: int = 4, 
                     axis: Union[int, str] = 0
                    ) -> Union[ndarray, DataArray]:
    """
    使用巴特沃斯带通滤波器滤波(scipy.signal.butter())
    参考自纪献普师兄的文章: 
    https://blog.csdn.net/weixin_44237337/article/details/124135064

    首先定义时间单位 T 为多长，例如 1 日或 1 年，以此为单位确定滤波的周期，
    最后以此为单位确定采样频率。
    
    例如过滤 1 日 4 次的某变量数据，滤波带宽为 2-10 天，令 T 为 1 日：
    butter_band_pass(data, 2, 10, fs=4)

    ----------------------------------
        
    :param data: 要进行滤波的数据，如果是 DataArray，则只替换数据值，保留其他属性。
    :type data: ndarray | DataArray
    :param Tmin: 最小周期，时间单位为T
    :type Tmin: float
    :param Tmax: 最大周期，时间单位为T
    :type Tmax: float
    :param N: 滤波器阶数
    :type N: int
    :param fs: 采样频率，一个时间单位 T 内的数据次数。例如令 T 为 1 天，数据是小时数据，
        则 fs = 24
    :param axis: 滤波的维度，如果是 DataArray，则可以使用维度名，如"time"。
    :type axis: int | str

    :return: 滤波后的数据
    :rtype: ndarray | DataArray

    """
    import scipy.signal as signal

    b, a = signal.butter(N, [2/(Tmax*fs), 2/(Tmin*fs)], 'bandpass')

    # 如果是xarray的数据，则只替换数据值，保留其他属性
    if isinstance(data, DataArray):
        if isinstance(axis, str):
            axis = data.dims.index(axis)
        filteredData = signal.filtfilt(b, a, data.values, axis=axis)
        data.values = filteredData
        return data
    
    else:
        if isinstance(axis, str):
            raise ValueError("ndarray 的 axis 参数不能为 str")

        filteredData = signal.filtfilt(b, a, data, axis=axis)
        return filteredData


def assign_units(variables: list[DataArray] | DataArray, unit: str) -> \
    list[DataArray] | DataArray:
    """
    为数据指定 metpy 单位，支持单个和多个数据同时指定单位。
    例如： u, v = assign_units([u, v], 'm/s')

    ----------------------------------

    :param variables:要指定单位的数据
    :type variables: list[DataArray] | DataArray
    :param unit: 单位，例如 'm/s'
    :type unit: str

    :return: 指定单位后的 metpy.xarray 数据，如果是多个数据，则返回 list。
    :rtype: list[DataArray] | DataArray
    """
    from metpy.units import units as mpunits

    if isinstance(variables, list):
        for var in variables:
            var = var * mpunits(unit)
    else:
        variables = variables * mpunits(unit)
    
    return variables


def vorticity(u: DataArray, v: DataArray, 
              lon: ndarray=None, lat: ndarray=None) -> DataArray:
    """
    计算涡度，metpy.calc.vorticity()的封装，传入带坐标的 DataArray，
    或者手动指定经纬度数组。 优先使用数据自带的经纬度坐标。

    ----------------------------------
    
    :param u: u风
    :type u: DataArray
    :param v: v风
    :type v: DataArray
    :param lon: optional 经度，一维数组
    :type lon: ndarray
    :param lat: optional 纬度，一维数组
    :type lat: ndarray

    :return: 涡度
    :rtype: DataArray

    """
    import metpy.calc as mpcalc

    if hasattr(u, 'longitude'):
        dx, dy = mpcalc.lat_lon_grid_deltas(u.longitude, u.latitude)
        vort = mpcalc.vorticity(u, v, dx=dx, dy=dy)

    elif hasattr(u, 'lon'):
        dx, dy = mpcalc.lat_lon_grid_deltas(u.lon, u.lat)
        vort = mpcalc.vorticity(u, v, dx=dx, dy=dy)
        
    elif lon is not None and lat is not None:
        dx, dy = mpcalc.lat_lon_grid_deltas(lon, lat)
        vort = mpcalc.vorticity(u, v, dx, dy)

    else:
        raise ValueError("无经纬度信息")

    return vort


def Kmean(data: Union[DataArray, ndarray], n_clusters: int=4, dim: str='time',
          silhouette: bool=False) -> Union[DataArray, ndarray, 
                                           tuple[DataArray, float], 
                                           tuple[ndarray, float]]:
    """计算K均值聚类的分类结果，传入 DataArray 可以指定聚类的维度，并返回带有此维度
        信息的聚类结果。可选择是否计算并返回聚类的轮廓系数。

    :param data: 要聚类的数据，如果是 DataArray，则会按照 dim 维度聚类；
        如果是 ndarray，则直接聚类，以第 0 维为样本列。
    :type data: Union[DataArray, ndarray]
    :param n_clusters: 分类数, defaults to 4
    :type n_clusters: int, optional
    :param dim: 聚类的维度名, defaults to 'time'
    :type dim: str, optional
    :param silhouette: 是否计算并返回聚类的轮廓系数, defaults to False
    :type silhouette: bool, optional
    :return: 聚类结果标签，如果数据是 DataArray，则返回带有聚类维度坐标的 DataArray
    :rtype: Union[DataArray, ndarray, tuple[DataArray, float], 
                    tuple[ndarray, float]]
    """
    from sklearn.cluster import KMeans

    if isinstance(data, DataArray):
        dimValues = data[dim]
        # 把这个维度转到第0
        data = data.transpose(
            *[dim] + [d for d in data.dims if d != dim]
        )
        data = data.values.reshape(data.shape[0], -1)
        kmeans = KMeans(n_clusters=n_clusters, n_init='auto').fit(data)
        labels = kmeans.labels_
        result = DataArray(labels, coords={dim: dimValues}, dims=[dim])

        if silhouette:
            from sklearn.metrics import silhouette_score
            return result, silhouette_score(data, labels)

        return result
    
    else:
        data = data.reshape(data.shape[0], -1)
        kmeans = KMeans(n_clusters=n_clusters, n_init='auto').fit(data)
        labels = kmeans.labels_

        if silhouette:
            from sklearn.metrics import silhouette_score
            return labels, silhouette_score(data, labels)

        return labels


def reverse_longitude(data: Union[DataArray, Dataset], 
                      lonName: str='longitude') -> Union[DataArray, Dataset]:
    """将数据的经度从 -180~180 转换为 0~360 度的形式

    :param data: xarray 数据
    :type data: Union[DataArray, Dataset]
    :param lonName: 经度维度的名称, defaults to 'longitude'
    :type lonName: str, optional
    :return: 转换后数据
    :rtype: Union[DataArray, Dataset]
    """
    param = {lonName: (data.coords[lonName].values + 360) % 360}
    data = data.assign_coords(**param)
    data = data.sortby(data.coords[lonName])

    return data


def lon_lat_distance(lon: Union[DataArray, ndarray], 
                     lat: Union[DataArray, ndarray]) -> \
                    tuple[ndarray[float, float], ndarray[float, float]]:
    """
    使用球模型，生成经纬度网格的距离，用于差分计算，如散度涡度等。
    使用CGCS2000的半长轴为球模型的半径。
    
    分辨率会根据经纬度前两个值的差分别计算，只支持等经纬度格点数据。

    ----------------------------------

    :param lon: 经度，一维数组
    :type lon: Union[DataArray, ndarray]
    :param lat: 纬度，一维数组
    :type lat: Union[DataArray, ndarray]

    :return: 返回 (dx, dy)，经纬度网格的距离，单位为 m，shape 为 (len(lat), len(lon))
    :rtype: tuple[ndarray[float, float], ndarray[float, float]]
    """

    if isinstance(lon, DataArray):
        lon = lon.values
    if isinstance(lat, DataArray):
        lat = lat.values

    resolutionX: float = np.abs(lon[1] - lon[0])
    resolutionY: float = np.abs(lat[1] - lat[0])

    dy = np.ones((len(lat), len(lon))) * const.LATITUDE_DISTANCE * resolutionY
    dx = np.ones((len(lat))) * const.LATITUDE_DISTANCE * \
        np.cos(lat * np.pi / 180) * resolutionX
    dx = np.expand_dims(dx, axis=1).repeat(len(lon), axis=1)
    return dx, dy


def grad(data: DataArray, dim:str, 
         lonName='longitude', latName='latitude') -> DataArray:
    """计算水平梯度，dim 参数指定梯度方向，是经度 lonName 或纬度 latName。
    例如计算散度：
        ```python
        pupx = grad(u, 'longitude')
        pypy = grad(v, 'latitude')
        div = pupx + pypy
        ```
        
    :param data: 要计算梯度的数据
    :type data: xr.DataArray
    :param dim: 梯度方向的维度名，只能与 lonName 或 latName 之一相同
    :type dim: str
    :param lonName: 经度维度名称, defaults to 'longitude'
    :type lonName: str, optional
    :param latName: 纬度维度名称, defaults to 'latitude'
    :type latName: str, optional
    :return: 梯度结果
    :rtype: DataArray
    """
    dx, dy = lon_lat_distance(data[lonName].values, data[latName].values)
    resolution = np.abs(data[lonName].values[1] - data[lonName].values[0])

    grad = data.differentiate(dim) * resolution # 乘上分辨率，得到差分
    grad = grad / (dx if dim == lonName else dy)
    return grad


def loop_moving_average(d: Union[ndarray, list], window: int = 3) -> ndarray:
    """
    循环数据的滑动平均，头尾相接滑动平均，滑动窗口为 `window * 2 + 1`。

    :param d: 一维数据
    :type d: Union[ndarray, list]
    :param window: 滑动窗口单侧大小
    :type window: int, optional

    :return: ndarray 滑动平均后的数据
    :rtype: ndarray
    """

    if isinstance(d, list):
        d = np.array(d)
    
    if len(d.shape) != 1:
        raise ValueError("只支持一维数据")

    # 循环滑动平均
    d = np.concatenate((d[-window:], d, d[:window]))
    dd = np.zeros_like(d)
    for i in range(window, len(d)-window):
        dd[i] = np.mean(d[i-window:i+window+1])
    dd = dd[window:-window]

    return dd


def mk_test(x: ndarray) -> tuple[ndarray, ndarray]:
    """
    MK检验,输入1维时间序列
    :param x: 一维时间序列
    :type x: ndarray

    :return: UF, UB
    :rtype: tuple[ndarray, ndarray]
    """

    def mk_sk(x: ndarray) -> ndarray:
        """
        MK检验的sk,输入1维时间序列
        :param x: 一维时间序列
        :type x: ndarray

        :return: ndarray
        """
        SK = [0]
        n = len(x)
        sk = 0
        E = [0]
        Var = [0]
        UF = [0]
        for i in range(1,n):
            for j in range(i):
                if x[j] < x[i]:
                    sk += 1
            SK.append(sk)
            E.append((i+0)*(i+1)/4)
            Var.append((i+1)*(i)*(2*(i+1)+5)/72)
            UF.append((SK[i]-E[i])/np.sqrt(Var[i]))
        UF = np.array(UF)
        return UF

    UF = mk_sk(x)
    UB = mk_sk(x[::-1])
    UB = -UB[::-1]
    return UF, UB


def crop_data(data: Union[DataArray, Dataset], 
              geoRange: Union[list[float], tuple[float], ndarray],
              lonName: str="longitude", latName: str="latitude") -> Union[DataArray, Dataset]:
    """根据经纬度裁剪数据，自动识别纬度顺序

    :param data: 源数据
    :type data: Union[DataArray, Dataset]
    :param geoRange: 地理范围，[lon_min, lon_max, lat_min, lat_max]
    :type geoRange: Union[list[float], tuple[float], ndarray]
    :param lonName: 经度维度名称, defaults to "longitude"
    :type lonName: str, optional
    :param latName: 纬度维度名称, defaults to "latitude"
    :type latName: str, optional
    :return: 裁剪后的数据
    :rtype: Union[DataArray, Dataset]
    """
    geoRange = geoRange.copy()

    # 检查 lat 顺序是从大到小还是从小到大
    if data[latName].values[0] > data[latName].values[-1]:
        geoRange[2], geoRange[3] = geoRange[3], geoRange[2]

    param = {lonName: slice(geoRange[0], geoRange[1]),
                latName: slice(geoRange[2], geoRange[3])}

    return data.sel(**param)


def interp_2d_data_in_time(data: DataArray,
                           z: ndarray, newZ: ndarray, 
                           method: str = 'linear', xName='longitude',
                           yName='latitude') -> ndarray:
    """
    用不同时间/情况的2d数据空间场，需要求其过程中的值。

    :param data: 2d 数据空间场，shape=(z, lat, lon)
    :type data: DataArray
    :param z: 2d 数据空间场的 z 坐标，可以是时间、位置等就是要插值的坐标
    :type z: ndarray
    :param newZ: 要插值的坐标，应该是 2 维的 n 行 3 列的 (n, (z, lat, lon))，
    如果是 3 维，则按照多个 2 维数据空间场插值，即 k*n*3 的 (k, n, (z, lat, lon))
    :type newZ: ndarray
    :param method: 插值方法, Supported are “linear”, “nearest”, “slinear”, 
        “cubic”, “quintic” and “pchip”. 
    :type method: str, optional

    :return: 插值后的数据，shape=(k, n)，如果newZ是2维的，则k=1
    :rtype: ndarray
    """
    from scipy.interpolate import RegularGridInterpolator
    x = data[xName].values
    y = data[yName].values
    interp = RegularGridInterpolator((z, y, x), data.values, method=method, 
                                     bounds_error=False, fill_value=None)
    result = []

    if len(newZ.shape) == 2:
        newZ = np.expand_dims(newZ, axis=0)
    for i in range(newZ.shape[0]):
        result.append(interp(newZ[i]))
    
    result = np.array(result)
    return result


def array_exclude(array: Union[ndarray, list],
                   which:Union[float, list, ndarray]) -> ndarray:
    """从数组中排除指定元素，只支持一维数组

    :param array: 输入数组或列表
    :type array: Union[ndarray, list]
    :param which: 要排除的值，单个或多个
    :type which: Union[float, ndarray]
    :return: 排除后的数组
    :rtype: ndarray
    """
    if isinstance(array, list):
        array = np.array(array)

    if type(which) == float or type(which) == int:
        which = np.array([which])

    return np.array([i for i in array if i not in which])